package encryption;
/*
 * This file is part of the AES and RSA in Java demo. It demonstrates
 * the combination of a symmetric and asymmetric encryption algorithm.
 * Copyright (C) 2005 - 2011 CodePlanet. All rights reserved.
 * For more informations, visit http://www.codeplanet.eu/.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */


import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.security.GeneralSecurityException;
import java.security.Key;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.SecureRandom;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;

/**
 * A basic class for encryption and decription using symmetric and asymmetric
 * algorithms.
 *
 * @author  CodePlanet
 * @version 1.0.1, 12/06/2010
 * @see     <a href="http://www.codeplanet.eu/">http://www.codeplanet.eu/</a>
 */
public class EncryptionWithKeyPair
{   
    /** The key size. */
    private static final int KEYSIZE = 1024;
    
    /** 
     * Private Constructor, all methods are static.
     */
    private EncryptionWithKeyPair()
    {
    }
    
    /**
     * Generate a RSA pair key (public.key, private.key) and stores
     * it in files.
     *
     * @param privateKeyFileName    name of the private key
     * @param publicKeyFileName     name of the public key
     */    
    public static void generateKey(String privateKeyFileName, String publicKeyFileName)
    {
        try {
            KeyPairGenerator pairgen = KeyPairGenerator.getInstance("RSA");
            SecureRandom random = new SecureRandom();
            pairgen.initialize(KEYSIZE, random);
            KeyPair keyPair = pairgen.generateKeyPair();
            ObjectOutputStream out = new ObjectOutputStream(new FileOutputStream(publicKeyFileName));
            out.writeObject(keyPair.getPublic());
            out.close();
            out = new ObjectOutputStream(new FileOutputStream(privateKeyFileName));
            out.writeObject(keyPair.getPrivate());
            out.close();
        } catch (IOException e) {
        	System.out.println(e.getMessage());
        } catch (GeneralSecurityException e) {
        	System.out.println(e.getMessage());
        }
    }

    /**
     * Encrypt a file with AES using the public RSA key.
     *
     * @param publicKeyFile     name of the public key
     * @param inputFile         name of the input file
     * @param outputFile        name of the output file
     */    
    public static void encrypt(String publicKeyFile, String inputFile, String outputFile)
    {
        try {
            KeyGenerator keygen = KeyGenerator.getInstance("AES");
            SecureRandom random = new SecureRandom();
            keygen.init(random);
            SecretKey key = keygen.generateKey();

            // Wrap with public key
            ObjectInputStream keyIn = new ObjectInputStream(new FileInputStream(publicKeyFile));
            Key publicKey = (Key) keyIn.readObject();
            keyIn.close();

            Cipher cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.WRAP_MODE, publicKey);
            byte[] wrappedKey = cipher.wrap(key);
            DataOutputStream out = new DataOutputStream(new FileOutputStream(outputFile));
            out.writeInt(wrappedKey.length);
            out.write(wrappedKey);

            InputStream in = new FileInputStream(inputFile);
            cipher = Cipher.getInstance("AES");
            cipher.init(Cipher.ENCRYPT_MODE, key);
            crypt(in, out, cipher);
            in.close();
            out.close();
        } catch (IOException e) {
        	System.out.println(e.getMessage());
        } catch (GeneralSecurityException e) {
        	System.out.println(e.getMessage());
        } catch (ClassNotFoundException e) {
        	System.out.println(e.getMessage());
        }        
    }
 
    /**
     * Decrypt a file with the private key.
     *
     * @param privateKeyFile
     * @param inputFile
     * @param outputFile
     */
    public static void decrypt(String privateKeyFile, String inputFile, String outputFile)
    {
        try {
            DataInputStream in = new DataInputStream(new FileInputStream(inputFile));
            int length = in.readInt();
            byte[] wrappedKey = new byte[length];
            in.read(wrappedKey, 0, length);

            // Open with private key
            ObjectInputStream keyIn = new ObjectInputStream(new FileInputStream(privateKeyFile));
            Key privateKey = (Key) keyIn.readObject();
            keyIn.close();

            Cipher cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.UNWRAP_MODE, privateKey);
            Key key = cipher.unwrap(wrappedKey, "AES", Cipher.SECRET_KEY);

            OutputStream out = new FileOutputStream(outputFile);
            cipher = Cipher.getInstance("AES");
            cipher.init(Cipher.DECRYPT_MODE, key);
            crypt(in, out, cipher);
            in.close();
            out.close();
        } catch (IOException e) {
        	System.out.println(e.getMessage());
        } catch (GeneralSecurityException e) {
        	System.out.println(e.getMessage());
        } catch (ClassNotFoundException e) {
        	System.out.println(e.getMessage());
        }
        
        try {
			
		} catch (Exception e) {
			// TODO: handle exception
		}
    }    

    /**
     * Transforms all bytes from the input stream with a crypt alorithm
     * and sends the transformed bytes to the output stream.
     * @param in        input stream
     * @param out       output stream
     * @param cipher    crypt alorithm used
     */
    public static void crypt(InputStream in, OutputStream out, Cipher cipher)
        throws IOException, GeneralSecurityException
    {
        int blockSize = cipher.getBlockSize();
        int outputSize = cipher.getOutputSize(blockSize);
        byte[] input = new byte[blockSize];
        byte[] output = new byte[outputSize];
        int inLength = 0;
        boolean finished = false;

        while (!finished) {
            inLength = in.read(input);

            if (inLength == blockSize) {
                int outLength = cipher.update(input, 0, blockSize, output);
                out.write(output, 0, outLength);
            } else {
                finished = true;
            }
        }

        if (inLength > 0)
            output = cipher.doFinal(input, 0, inLength);
        else
            output = cipher.doFinal();

        out.write(output);
    }

}
